{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 轨道旋转 MP2 方法 (OO-MP2) 简单理解"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 创建日期：2021-01-09"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这篇文档会尝试简单介绍轨道旋转 MP2 方法 (Orbital-Optimized Second-Order Møller–Plesset Perturbation, OO-MP2 or OMP2) 的基础概念与 PySCF 上的程序实现和理解。\n",
    "\n",
    "这篇文档的编写并没有翻阅很多文献，并作测评上的认识。为数不多的文献与参考资料是\n",
    "\n",
    "> Sun, Chan, et al. [^Sun-Chan.JCP.2020] (PySCF 进展文章)\n",
    ">   \n",
    "> PySCF 并没有一个完整或独立的 OO-MP2 模块。实现 OO-MP2 可以通过仿 CASSCF 的方式实现。之后使用到的 `MP2AsFCISolver` class 就是直接从该文章中截取的演示代码。\n",
    "\n",
    "> Psi4NumPy 演示文档 [10a_orbital-optimized-mp2.ipynb](https://github.com/psi4/psi4numpy/blob/master/Tutorials/10_Orbital_Optimized_Methods/10a_orbital-optimized-mp2.ipynb)\n",
    ">\n",
    ">   这是一份比较不错的基于 Psi4 的程序简要文档，使用的算法与技巧也不复杂。\n",
    "\n",
    "需要指出，这里的 OO-MP2 程序实现完全是基于 Post-HF 的闭壳层、无冻结轨道 MP2 实现的。更复杂的开壳层、冻结轨道、双杂化泛函方法，都不予以考虑。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy\n",
    "from pyscf import gto, mcscf, fci, mp, scf\n",
    "from functools import partial\n",
    "\n",
    "np.random.seed(0)\n",
    "np.einsum = partial(np.einsum, optimize=True)\n",
    "np.set_printoptions(precision=4, linewidth=150, suppress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这篇文档的程序理解部分，我们都会使用下述水分子。但文档末尾，我们会用氢分子的例子，说明 OO-MP2 的能量未必要比 MP2 能量要低。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyscf.gto.mole.Mole at 0x7efdd0fbc8b0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol = gto.Mole()\n",
    "mol.atom = \"\"\"\n",
    "O  0. 0. 0.\n",
    "H  0. 0. 1.\n",
    "H  0. 1. 0.\n",
    "\"\"\"\n",
    "mol.basis = \"6-31G\"\n",
    "mol.verbose = 0\n",
    "mol.build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PySCF 程序实现：高效方式"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这段程序 `MP2AsFCISolver` class 是直接从 Sun 的 JCP 文章截取的。通过在 CASSCF 中，将活性空间更改为全轨道、更改约化密度矩阵 (1-RDM, 2-RDM) 的生成方式为 MP2 的约化密度矩阵、并且允许活性空间的轨道旋转，就可以实现 OO-MP2。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MP2AsFCISolver:\n",
    "    def kernel(self, h1, h2, norb, nelec, ci0=None, ecore=0, **kwargs):\n",
    "        # Kernel takes the set of integrals from the current set of orbitals\n",
    "        fakemol = gto.M(verbose=0)\n",
    "        fakemol.nelectron = sum(nelec)\n",
    "        fake_hf = fakemol.RHF()\n",
    "        fake_hf._eri = h2\n",
    "        fake_hf.get_hcore = lambda *args: h1\n",
    "        fake_hf.get_ovlp = lambda *args: np.eye(norb)\n",
    "        \n",
    "        # Build an SCF object fake_hf without SCF iterations to perform MP2\n",
    "        fake_hf.mo_coeff = np.eye(norb)\n",
    "        fake_hf.mo_occ = np.zeros(norb)\n",
    "        fake_hf.mo_occ[:fakemol.nelectron//2] = 2\n",
    "        self.mp2 = fake_hf.MP2().run()\n",
    "        return self.mp2.e_tot + ecore, self.mp2.t2\n",
    "    \n",
    "    def make_rdm12(self, t2, norb, nelec):\n",
    "        dm1 = self.mp2.make_rdm1(t2)\n",
    "        dm2 = self.mp2.make_rdm2(t2)\n",
    "        return dm1, dm2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`mf_rhf` 为 RHF 实例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-75.9697009626036"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_rhf = mol.RHF().run()\n",
    "mf_rhf.e_tot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`mf_mp2` 为 MP2 实例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-76.1040356515777"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_mp2 = mp.MP2(mf_rhf).run()\n",
    "mf_mp2.e_tot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.13433468897410067"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_mp2.e_corr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`mf_cas` 在这里是指 OO-MP2 实例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-76.10510419427318"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_cas = mcscf.CASSCF(mf_rhf, mol.nao, mol.nelectron)\n",
    "mf_cas.fcisolver = MP2AsFCISolver()\n",
    "mf_cas.internal_rotation = True\n",
    "cas_result = mf_cas.kernel()\n",
    "cas_result[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PySCF 程序实现：大体思路拆解"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在这一段中，我们不会使用 PySCF 的 `CASSCF` class，而是从 RHF 与 MP2 的结果，了解 OO-MP2 的大体思路。\n",
    "\n",
    "从结果上，这种实现方式与 PySCF 会相同。但 PySCF 的 `CASSCF` class 一般会使用二阶方法 (即使用 Orbital Hessian) 加速收敛，而我们这里只使用一阶梯度下降方法 (Orbital Gradient) 进行收敛；一阶收敛方法显然会慢一些，但公式与程序会简单一些。\n",
    "\n",
    "首先对一些基础变量作声明：\n",
    "\n",
    "- `nocc` $n_\\mathrm{occ}$ 占据轨道数，`nvir` $n_\\mathrm{vir}$ 非占轨道数；\n",
    "\n",
    "- `nmo` $n_\\mathrm{MO}$ 分子轨道数，一般与原子轨道数相等；\n",
    "\n",
    "- `so` $[0:n_\\mathrm{occ}]$ 占据轨道分割 (slice)，`sv` $[n_\\mathrm{occ}:n_\\mathrm{MO}]$ 非占轨道分割 (slice)；\n",
    "\n",
    "- `mo_occ` PySCF 中用于表示轨道占据数的变量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "nocc, nmo = mol.nelec[0], mol.nao\n",
    "nvir = nmo - nocc\n",
    "so, sv = slice(0, nocc), slice(nocc, nmo)\n",
    "mo_occ = mf_rhf.mo_occ"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OO-MP2 的大体过程可以拆分为如下循环：\n",
    "\n",
    "1. 代入分子轨道系数 $C_{\\mu i}$，得到该系数下 MP2 的激发张量 $t_{ij}^{ab}$；\n",
    "\n",
    "2. 进而生成该情形下的 1-RDM $\\gamma_{pq}$ 与 2-RDM $\\Gamma_{pr}^{qs}$；\n",
    "\n",
    "3. 进而生成广义 Fock 矩阵 $F_{pq}$ 与轨道梯度 $x_{pq} = F_{pq} - F_{qp}$；\n",
    "\n",
    "4. 最后更新分子轨道系数 $C_{\\mu i}$。\n",
    "\n",
    "最终的收敛条件判据是 $F_{pq} - F_{qp} = 0$，即广义 Fock 矩阵 $F_{pq}$ 为对称矩阵。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def oomp2_cycle(C):\n",
    "    # Generate Psuedo objects, and therefore t_iajb\n",
    "    mf_prhf = scf.RHF(mol)\n",
    "    mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C\n",
    "    mf_pmp2 = mp.MP2(mf_prhf).run()                                                         # Step 1\n",
    "    # Generate 1-RDM, 2-RDM and orbital gradient from generalized Fock matrix\n",
    "    rdm1, rdm2 = mf_pmp2.make_rdm1(), mf_pmp2.make_rdm2()                                   # Step 2\n",
    "    gfock_grad = mf_cas.unpack_uniq_var(mf_cas.get_grad(C, (rdm1, rdm2), mf_cas.ao2mo(C)))  # Step 3\n",
    "    # Returned value: Updated MO Coefficient; OO-MP2 Energy (in current cycle); orbital gradient error\n",
    "    return update_C(C, gfock_grad), mf_pmp2.e_tot, np.linalg.norm(gfock_grad)               # Step 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "而更新轨道系数是通过下述过程实现：\n",
    "\n",
    "$$\n",
    "\\begin{gather}\n",
    "X_{ai} = - X_{ia} = \\frac{x_{ai}}{- \\varepsilon_a + \\varepsilon_i} = \\frac{F_{ai} - F_{ia}}{- \\varepsilon_a + \\varepsilon_i} \\\\\n",
    "X_{ij} = 0, \\; = X_{ab} = 0 \\\\\n",
    "\\mathbf{C} \\leftarrow \\mathbf{C} \\exp(\\lambda \\mathbf{X})\n",
    "\\end{gather}\n",
    "$$\n",
    "\n",
    "其中 $\\lambda$ 是梯度下降率，对应于机器学习，它与梯度下降算法的学习率是类似的。这里取为 $\\lambda = 0.5$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_C(C, gfock_grad):\n",
    "    # Generate anti-symmetrized rotation matrix\n",
    "    D = mf_rhf.make_rdm1(C, mo_occ)\n",
    "    e = (C.T @ mf_rhf.get_fock(dm=D) @ C).diagonal()\n",
    "    X = np.zeros_like(C)\n",
    "    X[sv, so] = gfock_grad[sv, so] / (- e[sv, None] + e[None, so])\n",
    "    X[so, sv] = gfock_grad[so, sv] / (- e[None, sv] + e[so, None])\n",
    "    # Control rotation by factor\n",
    "    X *= 0.5\n",
    "    # Generate rotated MO coefficient\n",
    "    C_new = C @ scipy.linalg.expm(X)\n",
    "    return C_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果将 RHF 的分子轨道系数 `mf_rhf.mo_coeff` 作为分子轨道系数的初猜，那么收敛过程可以用下述迭代代码给出："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "C_oo = np.copy(mf_rhf.mo_coeff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cycle | OO-MP2 Energy | G-Fock Gradient Norm\n",
      "    0 | -76.10403565  |  7.90255467e-02  \n",
      "    1 | -76.10503066  |  2.20049872e-02  \n",
      "    2 | -76.10509490  |  7.36831750e-03  \n",
      "    3 | -76.10510186  |  4.69833400e-03  \n",
      "    4 | -76.10510336  |  2.78455388e-03  \n",
      "    5 | -76.10510386  |  1.90302779e-03  \n",
      "    6 | -76.10510405  |  1.22381869e-03  \n",
      "    7 | -76.10510413  |  8.22563315e-04  \n",
      "    8 | -76.10510417  |  5.38178673e-04  \n",
      "    9 | -76.10510418  |  3.59079689e-04  \n",
      "   10 | -76.10510419  |  2.36428274e-04  \n",
      "   11 | -76.10510419  |  1.57182522e-04  \n",
      "   12 | -76.10510419  |  1.03794108e-04  \n",
      "   13 | -76.10510419  |  6.88759621e-05  \n",
      "   14 | -76.10510419  |  4.55466994e-05  \n"
     ]
    }
   ],
   "source": [
    "print(\"Cycle | OO-MP2 Energy | G-Fock Gradient Norm\")\n",
    "for i in range(15):\n",
    "    C_oo, eng, err = oomp2_cycle(C_oo)\n",
    "    print(\"{:5d} | {:<13.8f} |  {:<16.8e}\".format(i, eng, err))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{admonition} 记号区别\n",
    "\n",
    "在这份文档中，RHF 的 Fock 矩阵记号定义为 $f_{pq}$，而 Post-HF 方法的 Fock 矩阵记号定义为 $F_{pq}$。这两者并非相同，并且非轨道优化的方法下，广义 Fock 矩阵 $F_{pq}$ 矩阵一般是非对称的。\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PySCF 程序实现：理解与分解"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们会对上面程序中的重要步骤进行说明。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 原子轨道电子积分定义"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `h` $h_{\\mu \\nu}$，维度 $(\\mu, \\nu)$，原子轨道基的 Hamiltonian Core 矩阵，即动能与原子核-电子势能积分；\n",
    "\n",
    "- `S` $S_{\\mu \\nu}$，维度 $(\\mu, \\nu)$，原子轨道基的重叠积分；\n",
    "\n",
    "- `eri` $(\\mu \\nu | \\kappa \\lambda)$，维度 $(\\mu, \\nu, \\kappa, \\lambda)$，原子轨道基的双电子积分。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = mol.intor(\"int1e_kin\") + mol.intor(\"int1e_nuc\")\n",
    "S = mol.intor(\"int1e_ovlp\")\n",
    "eri = mol.intor(\"int2e\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Canonical MP2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们先简单回顾一下在 Canonical RHF 下，MP2 的激发系数 $t_{ij}^{ab}$ 与能量 $E_\\mathrm{corr}^\\mathsf{MP2}$ 的导出方式。我们留意到 PySCF 的自洽场过程给出的是 Canonical 情况，即分子轨道的 Fock 矩阵 $f_{pq}$ 是对角矩阵。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `C` $C_{\\mu p}$ 为分子轨道系数，`e` $e_p$ 为轨道能量；\n",
    "\n",
    "- `D_iajb` $D_{ij}^{ab}$ MP2 分母项，维度 $(i, a, j, b)$：\n",
    "\n",
    "    $$\n",
    "    D_{ij}^{ab} = \\varepsilon_i - \\varepsilon_a + \\varepsilon_j - \\varepsilon_b\n",
    "    $$\n",
    "\n",
    "- `eri_mo` $(pq|rs)$ 分子轨道基下的双电子积分，维度 $(p, q, r, s)$：\n",
    "\n",
    "    $$\n",
    "    (pq|rs) = C_{\\mu p} C_{\\nu q} (\\mu \\nu | \\kappa \\lambda) C_{\\kappa r} C_{\\lambda s}\n",
    "    $$\n",
    "\n",
    "- `t_iajb` $t_{ij}^{ab}$ MP2 激发系数：\n",
    "\n",
    "    $$\n",
    "    t_{ij}^{ab} = \\frac{(ia|jb)}{D_{ij}^{ab}}\n",
    "    $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "C, e = mf_rhf.mo_coeff, mf_rhf.mo_energy\n",
    "D_iajb = e[so, None, None, None] - e[None, sv, None, None] + e[None, None, so, None] - e[None, None, None, sv]\n",
    "eri_mo = np.einsum(\"up, vq, uvkl, kr, ls -> pqrs\", C, C, eri, C, C)\n",
    "t_iajb = eri_mo[so, sv, so, sv] / D_iajb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "因此，MP2 相关能可以写为 (参考值为 -0.134335 a.u.)\n",
    "\n",
    "$$\n",
    "E_\\mathrm{corr}^\\mathsf{MP2} = \\big( 2 t_{ij}^{ab} - t_{ij}^{ba} \\big) (ia|jb)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.13433468897410067"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "((2 * t_iajb - t_iajb.swapaxes(-1, -3)) * eri_mo[so, sv, so, sv]).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Non-Canonical MP2：PySCF 程序"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但对于 OO-MP2 而言，由于产生了轨道旋转，我们需要考察 Non-Canonical RHF 的 MP2。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Non-Canonical 意指 RHF 的 Fock 矩阵 $f_{pq}$ 是分块对角化的，即占据-非占和非占-占据分块 $f_{ia}$、$f_{ai}$ 均为零；而占据和非占分块 $f_{ij}$、$f_{ab}$ 的矩阵并非是对角化的。\n",
    "\n",
    "为了构造这样一个 Non-Canonical RHF 的情形，我们可以对 Canonical RHF 分子轨道系数矩阵 `C_rhf` 作下述变换，得到 Non-Canonical RHF 分子轨道系数矩阵 `C_rot`：\n",
    "\n",
    "$$\n",
    "\\mathbf{C} \\leftarrow \\mathbf{C} \\mathbf{U}\n",
    "$$\n",
    "\n",
    "上述的 $\\mathbf{U}$ 矩阵是分块对角化的正交矩阵。为了构造这样的正交矩阵，我们可以生成一个分块对角化、且反对称的 `X` $\\mathbf{X}$ 矩阵，并令 $\\mathbf{U} = \\exp(\\mathbf{X})$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "C_rhf = mf_rhf.mo_coeff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.random.randn(nmo, nmo)\n",
    "X[sv, so] = X[so, sv] = 0\n",
    "X -= X.T\n",
    "X *= 0.02\n",
    "C_rot = C_rhf @ scipy.linalg.expm(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由此构建出的 Non-Canonical 分子轨道 Fock 矩阵 $f_{pq}$ 是分块对角化的，即不一定要求 $f_{ij} = \\delta_{ij} \\varepsilon_i$ 与 $f_{ab} = \\delta_{ab} \\varepsilon_a$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-20.4748,  -0.0683,  -0.3396,  -1.0108,  -0.9682,  -0.    ,  -0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ,  -0.    ,   0.    ],\n",
       "       [ -0.0683,  -1.3485,  -0.0071,  -0.0425,  -0.0204,   0.    ,  -0.    ,   0.    ,   0.    ,   0.    ,   0.    ,   0.    ,  -0.    ],\n",
       "       [ -0.3396,  -0.0071,  -0.6668,  -0.0222,  -0.0172,  -0.    ,   0.    ,   0.    ,   0.    ,  -0.    ,  -0.    ,   0.    ,  -0.    ],\n",
       "       [ -1.0108,  -0.0425,  -0.0222,  -0.6362,  -0.0523,  -0.    ,   0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ,   0.    ,   0.    ],\n",
       "       [ -0.9682,  -0.0204,  -0.0172,  -0.0523,  -0.5539,  -0.    ,  -0.    ,  -0.    ,  -0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ],\n",
       "       [ -0.    ,   0.    ,  -0.    ,  -0.    ,  -0.    ,   0.1967,   0.0008,  -0.0182,   0.0501,  -0.002 ,   0.0269,  -0.0106,   0.0776],\n",
       "       [ -0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ,   0.0008,   0.2872,  -0.0032,   0.011 ,   0.0263,   0.0326,  -0.0324,   0.0407],\n",
       "       [ -0.    ,   0.    ,   0.    ,  -0.    ,  -0.    ,  -0.0182,  -0.0032,   0.9833,   0.0038,  -0.0097,   0.0066,   0.0088,  -0.0122],\n",
       "       [  0.    ,   0.    ,   0.    ,   0.    ,  -0.    ,   0.0501,   0.011 ,   0.0038,   1.1572,  -0.0006,  -0.0001,   0.0048,  -0.0293],\n",
       "       [  0.    ,   0.    ,  -0.    ,   0.    ,  -0.    ,  -0.002 ,   0.0263,  -0.0097,  -0.0006,   1.1616,  -0.0056,  -0.0042,   0.0036],\n",
       "       [ -0.    ,   0.    ,  -0.    ,  -0.    ,   0.    ,   0.0269,   0.0326,   0.0066,  -0.0001,  -0.0056,   1.2463,  -0.0016,  -0.0138],\n",
       "       [ -0.    ,   0.    ,   0.    ,   0.    ,   0.    ,  -0.0106,  -0.0324,   0.0088,   0.0048,  -0.0042,  -0.0016,   1.3518,  -0.0047],\n",
       "       [  0.    ,  -0.    ,  -0.    ,   0.    ,  -0.    ,   0.0776,   0.0407,  -0.0122,  -0.0293,   0.0036,  -0.0138,  -0.0047,   1.7381]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fock_rot = np.einsum(\"up, uv, vq -> pq\", C_rot, mf_rhf.get_fock(), C_rot)\n",
    "fock_rot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于这样的分子轨道系数矩阵 `C_rot`，PySCF 的程序照样可以给出正确的 MP2 相关能量 -0.134335 a.u. (其中 `mf_prhf` 是指虚假的 (Pseudo) RHF 实例)："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.13433467530628806"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_prhf = scf.RHF(mol)\n",
    "mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C_rot\n",
    "mf_pmp2 = mp.MP2(mf_prhf).run()\n",
    "mf_pmp2.e_corr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Non-Canonical MP2：激发系数 $t_{ij}^{ab}$ 迭代更新方式"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先为程序与公式说明，定义一些变量：\n",
    "\n",
    "- RHF Fock 对角线占据部分记为 `eo` $\\varepsilon_i = f_{ii}$；\n",
    "\n",
    "- RHF Fock 对角线非占部分记为 `ev` $\\varepsilon_a = f_{aa}$；\n",
    "\n",
    "- RHF Fock 去除对角线的占据分块记为 `fock_oo` $f'_{ij} = (1 - \\delta_{ij}) f_{ij}$；\n",
    "\n",
    "- RHF Fock 去除对角线的非占分块记为 `fock_vv` $f'_{ab} = (1 - \\delta_{ab}) f_{ab}$；\n",
    "\n",
    "- 双电子积分 `eri_mo` $(pq|rs)$；\n",
    "\n",
    "- 只包含占据-非占分块的双电子积分 `eri_iajb` $(ia|jb)$\n",
    "\n",
    "- MP2 分母 `D_iajb` $D_{ij}^{ab}$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "eo, ev = fock_rot.diagonal()[so], fock_rot.diagonal()[sv]\n",
    "fock_oo, fock_vv = fock_rot[so, so], fock_rot[sv, sv]\n",
    "fock_oop, fock_vvp = fock_oo - np.diag(eo), fock_vv - np.diag(ev)\n",
    "eri_mo = np.einsum(\"up, vq, uvkl, kr, ls -> pqrs\", C_rot, C_rot, eri, C_rot, C_rot)\n",
    "eri_iajb = eri_mo[so, sv, so, sv]\n",
    "D_iajb = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{caution}\n",
    "\n",
    "**变量重新定义**\n",
    "\n",
    "上面代码块中 `eo`, `ev`, `eri_mo`, `D_iajb` 就在 Non-Canonical 的系数矩阵 `C_rot` 下给出；但我们曾经也在 Canonical 系数矩阵下给出过类似的变量。\n",
    "\n",
    "由于我们会经常切换各种系数矩阵的旋转方式 (非旋转、Non-Canonical、Non-HF)，因此一些变量也会被复用与复写，也暂不区分旋转后与旋转前的分子轨道角标。这可能会对阅读造成困惑。\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "依据不同的微扰理论定义方式，Non-Canonical RHF 的 MP2 相关能可能与 Canonical RHF 的 MP2 相关能不同。因此这里采用两个相关能相同的定义。此时激发系数 $t_{ij}^{ab}$ 应当满足\n",
    "\n",
    "$$\n",
    "(ia|jb) = t_{kj}^{ab} f_{ki} + t_{ik}^{ab} f_{kj} - t_{ij}^{cb} f_{ca} - t_{ij}^{ac} f_{cb}\n",
    "$$\n",
    "\n",
    "上式是对等式右的 $k$ 进行求和的。如果现在用 $f_{ij} = f'_{ij} + \\delta_{ij} \\varepsilon_i$，$f_{ab} = f'_{ab} + \\delta_{ab} \\varepsilon_a$ 展开，那么上式可以写为\n",
    "\n",
    "$$\n",
    "(ia|jb) = t_{ij}^{ab} D_{ij}^{ab} + t_{kj}^{ab} f'_{ki} + t_{ik}^{ab} f'_{kj} - t_{ij}^{cb} f'_{ca} - t_{ij}^{ac} f'_{cb}\n",
    "$$\n",
    "\n",
    "整理上式，就可以得到迭代关系\n",
    "\n",
    "$$\n",
    "t_{ij}^{ab} \\leftarrow \\frac{(ia|jb) - t_{kj}^{ab} f'_{ki} - t_{ik}^{ab} f'_{kj} + t_{ij}^{cb} f'_{ca} + t_{ij}^{ac} f'_{cb}}{D_{ij}^{ab}}\n",
    "$$\n",
    "\n",
    "一般来说，如果轨道的旋转并不是很剧烈，那么 $f'_{ij}$, $f'_{ab}$ 两者的贡献较小，因此 $t_{ij}^{ab} \\simeq (ia|jb) / D_{ij}^{ab}$ 会是一个比较好的近似。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在此情形下，Non-Canonical MP2 的能量计算方式如下：\n",
    "\n",
    "$$\n",
    "E_\\mathrm{corr}^\\mathsf{MP2} = \\big( 2 t_{ij}^{ab} - t_{ij}^{ba} \\big) (ia|jb)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面的程序就是实现 Non-Canonical MP2 的流程。\n",
    "\n",
    "- `update_t_iajb` 即使用迭代关系，更新 $t_{ij}^{ab}$；\n",
    "\n",
    "- `calculate_noncan_mp2` 即计算 Non-Canonical MP2 相关能的函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def update_t_iajb(t_iajb):\n",
    "    t_iajb_new = np.zeros_like(t_iajb)\n",
    "    t_iajb_new += np.einsum(\"icjb, ca -> iajb\", t_iajb, fock_vvp)\n",
    "    t_iajb_new += np.einsum(\"iajc, cb -> iajb\", t_iajb, fock_vvp)\n",
    "    t_iajb_new -= np.einsum(\"iakb, kj -> iajb\", t_iajb, fock_oop)\n",
    "    t_iajb_new -= np.einsum(\"kajb, ki -> iajb\", t_iajb, fock_oop)\n",
    "    t_iajb_new += eri_iajb\n",
    "    t_iajb_new /= D_iajb\n",
    "    return t_iajb_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_noncan_mp2(t_iajb):\n",
    "    return ((2 * t_iajb - t_iajb.swapaxes(-1, -3)) * eri_iajb).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "随后声明初猜 $t_{ij}^{ab} \\simeq (ia|jb) / D_{ij}^{ab}$ 并以此迭代更新；并以 Canonical MP2 的相关能加以验证。在 5 次循环后，几乎收敛到了正确能量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_iajb = eri_mo[so, sv, so, sv] / D_iajb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error:   3.41981239e-03\n",
      "Error:   2.09994114e-03\n",
      "Error:   9.08474334e-05\n",
      "Error:   9.06169148e-05\n",
      "Error:   3.54576068e-06\n",
      "Error:   5.43397725e-06\n",
      "Error:   6.95752296e-08\n",
      "Error:   4.42378672e-07\n",
      "Error:  -2.06561550e-08\n",
      "Error:   4.46581002e-08\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    print(\"Error: {:16.8e}\".format(calculate_noncan_mp2(t_iajb) - mf_mp2.e_corr))\n",
    "    t_iajb = update_t_iajb(t_iajb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "事实上，在 PySCF 中，包含占据-非占轨道旋转的 Non-RHF 下的 MP2 方法，也是通过上述过程进行计算的。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MP2 1-RDM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于一阶约化密度 1-RDM $\\gamma_{pq}$，其需要通过分块的方式生成：\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\gamma_{ij}^\\mathsf{RHF} &= 2 \\delta_{ij} \\\\\n",
    "\\gamma_{ab}^\\mathsf{RHF} &= \\gamma_{ia}^\\mathsf{RHF} = \\gamma_{ai}^\\mathsf{RHF} = 0 \\\\\n",
    "\\gamma_{ij}^\\mathrm{corr} &= - 4 t_{ik}^{ab} t_{jk}^{ab} + 2 t_{ik}^{ba} t_{jk}^{ab} \\\\\n",
    "\\gamma_{ab}^\\mathrm{corr} &= 4 t_{ij}^{ac} t_{ij}^{bc} - 2 t_{ij}^{ca} t_{ij}^{bc} \\\\\n",
    "\\gamma_{ia}^\\mathrm{corr} &= \\gamma_{ai}^\\mathrm{corr} = 0 \\\\\n",
    "\\gamma_{pq} &= \\gamma_{pq}^\\mathsf{RHF} + \\gamma_{pq}^\\mathrm{corr}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "这种生成方式无关乎方法是否是 Canonical 的。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先生成 RHF 的 1-RDM `rdm1_rhf` $\\gamma_{pq}^\\mathsf{RHF}$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "rdm1_rhf = np.zeros((nmo, nmo))\n",
    "np.fill_diagonal(rdm1_rhf[so, so], 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "随后给出 MP2 相关部分所给出的 1-RDM `rdm1_corr` $\\gamma_{pq}^\\mathrm{corr}$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "rdm1_corr = np.zeros((nmo, nmo))\n",
    "rdm1_corr[so, so] = - 4 * np.einsum(\"iakb, jakb -> ij\", t_iajb, t_iajb) + 2 * np.einsum(\"ibka, jakb -> ij\", t_iajb, t_iajb)\n",
    "rdm1_corr[sv, sv] = 4 * np.einsum(\"iajc, ibjc -> ab\", t_iajb, t_iajb) - 2 * np.einsum(\"icja, ibjc -> ab\", t_iajb, t_iajb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "总 1-RDM `rdm1` $\\gamma_{pq}$ 可以通过简单相加获得："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rdm1 = rdm1_rhf + rdm1_corr\n",
    "np.allclose(rdm1, mf_pmp2.make_rdm1())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MP2 2-RDM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于二阶约化密度 2-RDM `rdm2` $\\Gamma_{pr}^{qs}$ (维度 $(p, q, r, s)$)，其也需要通过分块生成。首先生成 $\\Gamma_{ia}^{jb}$, $\\Gamma_{ai}^{bj}$, $\\Gamma_{ik}^{jl}$, $\\Gamma_{ac}^{bd}$ 部分：\n",
    "\n",
    "$$\n",
    "\\Gamma_{pr}^{qs} = \\left( \\gamma_{pq} \\gamma_{rs} - \\frac{1}{2} \\gamma_{ps} \\gamma_{rq} \\right) - \\left( \\gamma_{pq}^\\mathrm{corr} \\gamma_{rs}^\\mathrm{corr} - \\frac{1}{2} \\gamma_{ps}^\\mathrm{corr} \\gamma_{rq}^\\mathrm{corr} \\right)\n",
    "$$\n",
    "\n",
    "其余的部分是 $\\Gamma_{ij}^{ab}$ 与 $\\Gamma_{ab}^{ij}$：\n",
    "\n",
    "$$\n",
    "\\Gamma_{ij}^{ab} = \\Gamma_{ab}^{ij} = 4 t_{ij}^{ab} - 2 t_{ij}^{ba}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rdm2 = np.zeros((nmo, nmo, nmo, nmo))\n",
    "rdm2 = np.einsum(\"pq, rs -> pqrs\", rdm1, rdm1) - 0.5 * np.einsum(\"ps, rq -> pqrs\", rdm1, rdm1)\n",
    "rdm2 -= np.einsum(\"pq, rs -> pqrs\", rdm1_corr, rdm1_corr) - 0.5 * np.einsum(\"ps, rq -> pqrs\", rdm1_corr, rdm1_corr)\n",
    "rdm2[so, sv, so, sv] = 4 * np.einsum(\"iajb -> iajb\", t_iajb) - 2 * np.einsum(\"ibja -> iajb\", t_iajb)\n",
    "rdm2[sv, so, sv, so] = 4 * np.einsum(\"iajb -> aibj\", t_iajb) - 2 * np.einsum(\"ibja -> aibj\", t_iajb)\n",
    "np.allclose(rdm2, mf_pmp2.make_rdm2(), atol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由此，我们可以通过 1-RDM $\\gamma_{pq}$ 与 2-RDM $\\Gamma_{pr}^{qs}$ 验证 MP2 总能量 -76.104036 a.u.：\n",
    "\n",
    "$$\n",
    "E_\\mathrm{tot}^\\mathsf{MP2} = h_{pq} \\gamma_{pq} + \\frac{1}{2} (pq|rs) \\Gamma_{pr}^{qs} + E_\\mathrm{nuc}\n",
    "$$\n",
    "\n",
    "但这里的单电子积分 $h_{pq}$ 与双电子积分 $(pq|rs)$ 都是在旋转过后的系数轨道矩阵 `C_rot` $\\mathbf{C}$ 为基给出，因此需要重新生成一下。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_mo = np.einsum(\"up, uv, vq -> pq\", C_rot, h, C_rot)\n",
    "eri_mo = np.einsum(\"up, vq, uvkl, kr, ls -> pqrs\", C_rot, C_rot, eri, C_rot, C_rot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-76.10403565383504"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    + np.einsum(\"pq, pq ->\", h_mo, rdm1)\n",
    "    + 0.5 * np.einsum(\"pqrs, pqrs ->\", eri_mo, rdm2)\n",
    "    + mol.energy_nuc()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成广义 Fock 矩阵"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "广义 Fock 矩阵 `gfock` $F_{pq}$ 区别于 RHF 的 Fock 矩阵 $f_{pq}$。其定义为\n",
    "\n",
    "$$\n",
    "F_{pq} = h_{pm} \\gamma_{mq} + (pm|rs) \\Gamma_{mr}^{qs}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfock = np.einsum(\"pr, rq -> pq\", h_mo, rdm1) + np.einsum(\"pmrs, mqrs -> pq\", eri_mo, rdm2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "事实上，RHF 的 Fock 矩阵中，占据轨道部分也可以用类似的方法定义：\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "2 f_{ij} &= h_{im} \\gamma_{mj}^\\mathsf{RHF} + (im|rs) \\Gamma_{mr}^{js, \\mathsf{RHF}} \\\\\n",
    "\\Gamma_{pr}^{qs, \\mathsf{RHF}} &= \\gamma_{pq}^\\mathsf{RHF} \\gamma_{rs}^\\mathsf{RHF} - \\frac{1}{2} \\gamma_{ps}^\\mathsf{RHF} \\gamma_{rq}^\\mathsf{RHF}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rdm2_rhf = np.einsum(\"pq, rs -> pqrs\", rdm1_rhf, rdm1_rhf) - 0.5 * np.einsum(\"ps, rq -> pqrs\", rdm1_rhf, rdm1_rhf)\n",
    "np.allclose(\n",
    "    (np.einsum(\"pr, rq -> pq\", h_mo, rdm1_rhf) + np.einsum(\"pmrs, mqrs -> pq\", eri_mo, rdm2_rhf))[so, so],\n",
    "    2 * fock_rot[so, so],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但在 PySCF 的 CASSCF 模块中，似乎没有直接生成广义 Fock 矩阵的方式。但其有广义 Fock 的导数量，被称为轨道梯度 (Orbital Gradient) `gfock_grad` $x_{pq}$：\n",
    "\n",
    "$$\n",
    "x_{pq} = F_{pq} - F_{qp}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gfock_grad = gfock - gfock.T\n",
    "np.allclose(\n",
    "    mf_cas.unpack_uniq_var(mf_cas.get_grad(C_rot, (rdm1, rdm2), mf_cas.ao2mo(C_rot))),\n",
    "    gfock_grad\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "至此，所有生成 OO-MP2 所需要的单步复杂计算都已经涵盖到了。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 轨道旋转的意义"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "讨论到现在，我们仅仅知道了 OO-MP2 的程序实现是如何进行的；但对其根源的合理性问题，我们在这里才开始说明。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "出于一般性，我们现在考虑 Non-HF 形式的轨道系数，即相对于 RHF 系数已经一定程度的旋转。该 Non-HF 轨道系数称为 `C_base` $C_{\\mu p}^\\mathrm{base}$。我们之后的讨论都基于该 Non-HF 轨道系数开始。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.random.randn(nmo, nmo)\n",
    "X = (X - X.T) * 0.02\n",
    "C_base = C_rhf @ scipy.linalg.expm(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先需要说明，轨道的旋转矩阵必须是正交矩阵 (酉矩阵)。这是因为轨道系数必须满足\n",
    "\n",
    "$$\n",
    "\\mathbf{C}^\\dagger \\mathbf{S} \\mathbf{C} = \\mathbf{I}\n",
    "$$\n",
    "\n",
    "旋转矩阵 $\\mathbf{U}$ 通过下式定义：$\\mathbf{C} = \\mathbf{C}^\\mathrm{base} \\mathbf{U}$。因此，\n",
    "\n",
    "$$\n",
    "\\mathbf{C}^\\dagger \\mathbf{S} \\mathbf{C} = \\mathbf{U}^\\dagger \\mathbf{C}^\\dagger \\mathbf{S} \\mathbf{C} \\mathbf{U} = \\mathbf{U}^\\dagger \\mathbf{I} \\mathbf{U} = \\mathbf{U}^\\dagger \\mathbf{U} = \\mathbf{I}\n",
    "$$\n",
    "\n",
    "而任何正交矩阵都可以通过反对称矩阵 $\\mathbf{X} = \\mathbf{X}^\\dagger$ 的幂次给出 $\\mathbf{U} = \\exp(\\mathbf{X})$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在考察在微扰下，能量随轨道系数的变化情况。令一般情况下轨道系数 $C_{\\mu p}$ 为关于反对称矩阵 $X_{pq}$ 的函数：\n",
    "\n",
    "$$\n",
    "\\mathbf{C} = \\mathbf{C}^\\mathrm{base} \\exp (\\mathbf{X})\n",
    "$$\n",
    "\n",
    "而 $C_{\\mu p}$ 对应的 MP2 能量写作关于 $X_{pq}$ 的函数 $E_\\mathrm{tot}^\\mathsf{MP2} (\\mathbf{X})$。下面的代码 `eng_mp2_pert` 即是代入反对称矩阵 $X_{pq}$，生成 MP2 能量的函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eng_mp2_pert(X):\n",
    "    C_rot = C_base @ scipy.linalg.expm(X)\n",
    "    mf_prhf = scf.RHF(mol)\n",
    "    mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C_rot\n",
    "    mf_pmp2 = mp.MP2(mf_prhf).run()\n",
    "    return mf_pmp2.e_tot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由此，能量关于旋转矩阵的导数关系可以写为矩阵 `dX` ${\\mathrm{d} \\mathbf{X}}$，其维度为 $(p, q)$：\n",
    "\n",
    "$$\n",
    "{\\mathrm{d}X}_{pq} = \\frac{\\partial E_\\mathrm{tot}^\\mathsf{MP2}}{\\partial X_{pq}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这种导数可以写成三点差分的数值微分的形式：\n",
    "\n",
    "$$\n",
    "{\\mathrm{d}X}_{pq} \\simeq \\frac{E_\\mathrm{tot}^\\mathsf{MP2} (d_{pq}) - E_\\mathrm{tot}^\\mathsf{MP2} (- d_{pq})}{2 d_{pq}}\n",
    "$$\n",
    "\n",
    "$E_\\mathrm{tot}^\\mathsf{MP2} (d_{pq})$ 的意义是，反对称矩阵 $\\mathbf{X}$ 仅在第 $p$ 行、第 $q$ 列上，$X_{pq} = d_{pq}$；且在第 $q$ 行、第 $p$ 列上，$X_{qp} = - d_{pq}$；其它位置上，$\\mathbf{X}$ 均取到零值。如果 $p = q$，那么 $\\mathbf{X} = \\mathbf{0}$。生成这种反对称矩阵的函数 `gen_pert_X` 如下所示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_pert_X(p, q, interval):\n",
    "    X = np.zeros((nmo, nmo))\n",
    "    X[p, q] = interval\n",
    "    X -= X.T\n",
    "    return X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "那么依据上述反对称矩阵，所求出的 MP2 能量随 $X_{pq}$ 变化的数值导数 ${\\mathrm{d}X}_{pq}$ 的生成函数如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eng_mp2_numdiff(p, q, interval):\n",
    "    X_positive = gen_pert_X(p, q, interval)\n",
    "    X_negative = gen_pert_X(p, q, -interval)\n",
    "    return (eng_mp2_pert(X_positive) - eng_mp2_pert(X_negative)) / (2 * interval)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对角标 $p, q$ 循环，我们就能求出完整的导数矩阵 `dX` ${\\mathrm{d} \\mathbf{X}}$ (这里选取的数值微分的间隙值 `interval` 为 $10^{-4}$ a.u.)："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.    , -0.    , -0.    , -0.    ,  0.    ,  0.243 , -0.6191,  0.7465, -1.7459,  1.0327, -0.5983,  1.3028, -1.8584],\n",
       "       [ 0.    ,  0.    , -0.    ,  0.    ,  0.    , -0.044 , -0.0219,  0.3123, -0.1693, -0.1755,  0.1931, -0.0509,  0.033 ],\n",
       "       [ 0.    ,  0.    ,  0.    , -0.    ,  0.    , -0.0894,  0.0924,  0.2349,  0.1175,  0.0443, -0.3922,  0.1505, -0.4868],\n",
       "       [ 0.    , -0.    ,  0.    ,  0.    , -0.    ,  0.0648, -0.0899, -0.0568, -0.0668, -0.137 ,  0.2291, -0.0017, -0.1029],\n",
       "       [-0.    , -0.    , -0.    ,  0.    ,  0.    , -0.0091,  0.0252,  0.1021, -0.0093, -0.0796,  0.0327, -0.067 , -0.0761],\n",
       "       [-0.243 ,  0.044 ,  0.0894, -0.0648,  0.0091,  0.    , -0.    ,  0.    , -0.    ,  0.    ,  0.    ,  0.    , -0.    ],\n",
       "       [ 0.6191,  0.0219, -0.0924,  0.0899, -0.0252,  0.    ,  0.    ,  0.    , -0.    ,  0.    ,  0.    ,  0.    , -0.    ],\n",
       "       [-0.7465, -0.3123, -0.2349,  0.0568, -0.1021, -0.    , -0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    , -0.    ],\n",
       "       [ 1.7459,  0.1693, -0.1175,  0.0668,  0.0093,  0.    ,  0.    , -0.    ,  0.    ,  0.    ,  0.    ,  0.    , -0.    ],\n",
       "       [-1.0327,  0.1755, -0.0443,  0.137 ,  0.0796, -0.    , -0.    , -0.    , -0.    ,  0.    , -0.    ,  0.    ,  0.    ],\n",
       "       [ 0.5983, -0.1931,  0.3922, -0.2291, -0.0327, -0.    , -0.    , -0.    , -0.    ,  0.    ,  0.    , -0.    ,  0.    ],\n",
       "       [-1.3028,  0.0509, -0.1505,  0.0017,  0.067 , -0.    , -0.    , -0.    , -0.    , -0.    ,  0.    ,  0.    , -0.    ],\n",
       "       [ 1.8584, -0.033 ,  0.4868,  0.1029,  0.0761,  0.    ,  0.    ,  0.    ,  0.    , -0.    , -0.    ,  0.    ,  0.    ]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dX = np.zeros((nmo, nmo))\n",
    "for a in range(nmo):\n",
    "    for i in range(nmo):\n",
    "        dX[a, i] = eng_mp2_numdiff(a, i, 1e-4)\n",
    "dX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "注意到这是一个反对称且分块的矩阵；在占据与非占分块值完全为零，有值处仅有 $\\mathrm{d} X_{ai} = - \\mathrm{d} X_{ia}$。这实际上近乎等于 2 倍的轨道梯度矩阵 `2 * gfock_grad`：\n",
    "\n",
    "$$\n",
    "\\mathrm{d} X_{pq} = 2 x_{pq} = 2 (F_{pq} - F_{qp})\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "mf_prhf = scf.RHF(mol)\n",
    "mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C_base\n",
    "mf_pmp2 = mp.MP2(mf_prhf).run()\n",
    "rdm1, rdm2 = mf_pmp2.make_rdm1(), mf_pmp2.make_rdm2()\n",
    "gfock_grad = mf_cas.unpack_uniq_var(mf_cas.get_grad(C_base, (rdm1, rdm2), mf_cas.ao2mo(C_base)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(2 * gfock_grad, dX, atol=5e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "因此，可以说 OO-MP2 的意义是，找到一个合适的 $\\mathbf{C}^\\mathrm{base}$，使得对于任意的很小的、用于旋转的反对称矩阵 $\\mathbf{X}$，有 $E_\\mathrm{tot}^\\mathsf{MP2} (\\mathbf{X})$ 不会更改。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OO-MP2 能量并非一定比 MP2 低"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在文档最后，我们会指出，OO-MP2 能量并非 MP2 的下界。尽管 OO-MP2 看起来对轨道进行变分式的优化，但其变分的对象应当认为是 Hylleraas 泛函，而非总 MP2 能量。\n",
    "\n",
    "对于下述拉长的氢分子，就是一个 OO-MP2 能量比 MP2 能量高的例子。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyscf.gto.mole.Mole at 0x7efd968f4310>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol = gto.Mole()\n",
    "mol.atom = \"\"\"\n",
    "H  0. 0. 0.\n",
    "H  0. 0. 15.\n",
    "\"\"\"\n",
    "mol.basis = \"6-31G\"\n",
    "mol.verbose = 0\n",
    "mol.build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "其 MP2 能量为"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-1.7458592201255043"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol.RHF().run().MP2().run().e_tot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "而其 OO-MP2 能量为"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-1.7280760742391805"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_cas = mcscf.CASSCF(mol.RHF().run(), mol.nao, mol.nelectron)\n",
    "mf_cas.fcisolver = MP2AsFCISolver()\n",
    "mf_cas.internal_rotation = True\n",
    "cas_result = mf_cas.kernel()\n",
    "cas_result[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但即使 OO-MP2 的能量比 MP2 高，它仍然无法解决 MP2 方法在解离两个氢原子所产生的相当大的解离误差。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[^Sun-Chan.JCP.2020]: Recent Developments in the PySCF Program Package. *J. Chem. Phys.* **2020**, *153* (2), 24109. doi: [10.1063/5.0006074](https://doi.org/10.1063/5.0006074)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
